import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import yfinance as yf

# 策略参数设置
initial_capital = 1000  # 初始本金
leverage = 20  # 杠杆倍数
grid_size = 50  # 网格间距（点数）
num_grids = 20  # 网格层数
transaction_cost = 0.0005  # 交易手续费率

# 生成示例价格数据（比特币永续合约）
# np.random.seed(42)
# dates = pd.date_range("2023-01-01", periods=500)
# prices = np.cumprod(1 + np.random.normal(0, 0.015, len(dates))) * 45000
# prices = pd.Series(prices, index=dates)

# 获取股票信息
stock = yf.Ticker("BTC-USD")
history = stock.history(period="15d")
prices = history["Close"].squeeze()


class NeutralGridStrategy:
    def __init__(self):
        self.equity = [initial_capital]  # 账户净值记录
        self.position = 0  # 净头寸
        self.entry_prices = []  # 持仓成本价记录
        self.grid_levels = []  # 网格价格层级
        self.trade_log = []  # 交易记录

    def initialize_grids(self, initial_price):
        # 创建上下对称的网格层级
        self.grid_levels = np.arange(
            initial_price - num_grids * grid_size,
            initial_price + (num_grids + 1) * grid_size,
            grid_size,
        )

    def execute_strategy(self, price_series):
        current_price = price_series[0]
        self.initialize_grids(current_price)

        for price in price_series:
            # 计算当前价格所在网格位置
            grid_idx = np.searchsorted(self.grid_levels, price)

            # 中性调整逻辑
            if grid_idx < len(self.grid_levels) // 2:
                # 价格低于中轨，执行做多
                self.execute_trade(price, "long")
            elif grid_idx > len(self.grid_levels) // 2:
                # 价格高于中轨，执行做空
                self.execute_trade(price, "short")
            else:
                # 中轨附近保持中性
                self.close_positions(price)

            # 更新账户净值
            self.update_equity(price)

    def execute_trade(self, price, direction):
        # 计算合约数量（含杠杆）
        contract_size = (initial_capital * leverage) / price
        cost = contract_size * price * transaction_cost

        if direction == "long":
            self.position += contract_size
            self.entry_prices.append(price)
            self.trade_log.append(("BUY", price, contract_size))
            self.equity[-1] -= cost
        elif direction == "short":
            self.position -= contract_size
            self.entry_prices.append(price)
            self.trade_log.append(("SELL", price, contract_size))
            self.equity[-1] -= cost

    def close_positions(self, price):
        if self.position != 0:
            # 平仓所有头寸
            realized_pnl = self.position * (price - np.mean(self.entry_prices))
            self.equity[-1] += realized_pnl
            self.position = 0
            self.entry_prices = []

    def update_equity(self, price):
        # 计算未实现盈亏
        unrealized_pnl = 0
        if self.position != 0 and self.entry_prices:
            avg_price = np.mean(self.entry_prices)
            unrealized_pnl = self.position * (price - avg_price)

        # 更新净值（含杠杆效应）
        self.equity.append(self.equity[-1] + unrealized_pnl)


# 实例化并运行策略
strategy = NeutralGridStrategy()
strategy.execute_strategy(prices)

# 可视化回测结果
plt.figure(figsize=(14, 10))

# 价格和网格可视化
ax1 = plt.subplot(3, 1, 1)
plt.plot(prices, label="Price", color="#1f77b4")
for level in strategy.grid_levels:
    plt.axhline(level, color="gray", linestyle="--", alpha=0.3)
plt.axhline(np.median(strategy.grid_levels), color="red", label="Median Grid")
plt.title("Price Movement with Grid Levels")
plt.ylabel("Price")
plt.legend()

# 头寸可视化
ax2 = plt.subplot(3, 1, 2, sharex=ax1)
position_changes = [t[2] if t[0] == "BUY" else -t[2] for t in strategy.trade_log]
plt.bar(
    prices.index[len(prices) - len(position_changes) :],
    position_changes,
    color=np.where(np.array(position_changes) > 0, "g", "r"),
    alpha=0.7,
)
plt.title("Position Changes")
plt.ylabel("Contract Size")

# 资金曲线可视化
ax3 = plt.subplot(3, 1, 3, sharex=ax1)
plt.plot(prices.index, strategy.equity[:-1], label="Account Equity", color="#2ca02c")
plt.fill_between(prices.index, strategy.equity[:-1], alpha=0.2, color="#2ca02c")
plt.title("Equity Curve")
plt.ylabel("Equity Value")

plt.tight_layout()
plt.show()

# 输出绩效指标
final_return = (strategy.equity[-2] / initial_capital - 1) * 100
max_drawdown = (
    (np.max(strategy.equity) - np.min(strategy.equity)) / np.max(strategy.equity) * 100
)
print(f"策略总收益: {final_return:.2f}%")
print(f"最大回撤: {max_drawdown:.2f}%")
print(f"总交易次数: {len(strategy.trade_log)}")
print(f"最终账户净值: {strategy.equity[-2]:.2f}")
